"""
    @brief:simulation environment
    @Editor:CJH
    @Date:2025/4/24
"""
import os.path
import threading
import time
import sys

sys.path.append("utils")

from dataloader import Dataloader

import cv2
import numpy as np
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from collections import deque

font = cv2.FONT_HERSHEY_COMPLEX_SMALL

start = time.time()


# 机器人对象
class BOT:
    def __init__(self, name: str, start_point: tuple, total_stride: int):
        self.name = name
        self.start_point = start_point  # np.array(x,y)
        self.total_stride = total_stride


# 设置障碍层，障碍的形状为rectangle grid
class Obstacle:
    def __init__(self, H, W):
        self.H = H
        self.W = W

    @staticmethod
    def clamp(point, H, W):
        return np.clip(np.array(point), [0, 0], [W - 1, H - 1])  # Return the clamped point as a tuple

    def add_Obstacle(self, start_point: tuple, end_point: tuple, input_map: np.ndarray):
        start_point = self.clamp(start_point, self.H, self.W)  # point(y,x)
        end_point = self.clamp(end_point, self.H, self.W)
        x1, y1 = np.minimum(start_point, end_point)
        x2, y2 = np.maximum(start_point, end_point)
        input_map[y1:y2 + 1, x1:x2 + 1] = 1
        return input_map


class Colors:
    def __init__(self):
        # hex = matplotlib.colors.TABLEAU_COLORS.values()
        hexs = ('344593', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
                '2C99A8', '00C2FF', 'FF3838', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
        self.n = len(self.palette)

    def rgb_call(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))

    @staticmethod
    def rgb2hex(rgb):
        """RGB转HEX
        :param rgb: RGB颜色元组，Tuple[int, int, int]
        :return: int or str
        """
        r, g, b = rgb
        result = r + (g << 8) + (b << 16)
        return hex(result)[2:]


# 只作为构建一个动态地图环境，可随时更新，用于训练DQN模型
class Map(Obstacle):
    def __init__(self, canvas_size=(800, 800), split_x_nums=20, split_y_nums=20,
                 obstacle_list=None):  # split_nums将地图分割为不同大小得栅格地图
        # 定义地图相关信息
        super(Map, self).__init__(split_y_nums, split_x_nums)
        self.params = None
        self.obstacle_map = None
        self.obstacle_map_normalized = None
        self.title = "Rectangle board"
        # 卷积核，使机器人移动一步的时候，有周围一圈的结果
        self.kernel = np.ones((3, 3), np.uint8)
        self.canvas_shape = (*canvas_size, 3)  # shape=(y,x)
        self.grid_map = np.zeros((split_y_nums, split_x_nums))  # shape=(y,x)
        # 初始化障碍层
        self.obstacle_list = obstacle_list
        # 初始化气体分数地图
        self.mass_fraction_viewer = Dataloader(self.grid_map.shape)
        # Declare all maps
        self.canvas, self.mass_fraction_map, self.Gaussian_mass_map = None, None, None
        # Define elements present inside the environment
        # 构造地图边界，确保机器人不会越界
        self.y_min = 0
        self.x_min = 0
        self.y_max = self.canvas_shape[0]
        self.x_max = self.canvas_shape[1]

    # 利用有限差分，计算障碍物的影响
    def update_diffusion(self, current_map, dt=1, D=0.1):
        new_map = current_map.copy()
        for y in range(1, self.H - 1):
            for x in range(1, self.W - 1):
                if self.obstacle_map[y, x] == 0:
                    laplacian = (current_map[y + 1, x] + current_map[y - 1, x] +
                                 current_map[y, x + 1] + current_map[y, x - 1] -
                                 4 * current_map[y, x])
                    new_map[y, x] += D * dt * laplacian
        return np.clip(new_map, 0, None) * (1 - self.obstacle_map)

    def fit_map(self, view=True):
        # 设置各向异性的法雷尔烟羽扩散方程,相当于曲面拟合
        def farrell_concentration(xy_data, x0, y0, sigma):
            x = xy_data[:, 0]
            y = xy_data[:, 1]
            d = np.sqrt((x - x0) ** 2 + (y - y0) ** 2)
            return np.max(self.mass_fraction_map) * np.exp(- (d ** 2) / (2 * (sigma ** 2)))

        # 1.可观测地图
        dilated_path = cv2.dilate(self.grid_map, self.kernel, iterations=1)
        Known_scores_map = np.array(dilated_path * self.mass_fraction_map)
        # 2.利用这个可观测地图的信息，进行拟合
        y_data, x_data = np.where(Known_scores_map > 0)
        c_data = self.mass_fraction_map[y_data, x_data]
        fractions_max_pos_ = np.unravel_index(np.argmax(Known_scores_map), Known_scores_map.shape)  # y,x

        p0 = [fractions_max_pos_[1], fractions_max_pos_[0], np.std(np.sqrt(y_data ** 2 + x_data ** 2))]
        xy_data = np.column_stack((x_data, y_data))  # 组合作为输入，求解参数
        params, _ = curve_fit(farrell_concentration, xy_data, c_data, p0=p0)

        # 由于在浓度区域对其扩散系数的拟合很差，添加高斯噪声扰动，使其脱离出低浓度地区
        params[-1] = params[-1] * np.exp(
            np.random.normal(0, 1 - np.max(Known_scores_map) / np.max(self.mass_fraction_map) + 1e-8))
        x0, y0, sigma = params
        # 3.使用拟合的参数，输出拟合扩散分布地图
        y_data, x_data = np.where(Known_scores_map >= 0)
        xy_data = np.column_stack((x_data, y_data))  # 组合作为输入，求解参数
        predict_mass_scores = farrell_concentration(xy_data, x0, y0, sigma)
        predict_mass_scores_map = np.zeros_like(self.grid_map)

        predict_mass_scores_map[xy_data[:, 1], xy_data[:, 0]] = predict_mass_scores * np.exp(
            np.random.normal(0, 0.7, size=len(xy_data)))

        # 4.考虑边界效应，利用有限差分计算障碍层的边界效应
        predict_mass_scores_map = self.update_diffusion(current_map=predict_mass_scores_map)
        Known_scores_map = np.where(Known_scores_map == 0, -0.1, Known_scores_map)
        if view:
            plt.imshow(predict_mass_scores_map / np.max(predict_mass_scores_map) - Known_scores_map, cmap='jet',
                       vmin=0)  # cmap 指定颜色映射
            plt.title('fit predict Heatmap Example')  # 设置标题
            plt.colorbar(label='Value')  # 添加颜色条
            plt.xlabel('X-axis')  # X 轴标签
            plt.ylabel('Y-axis')  # Y 轴标签
            plt.show()

        return self.clamp((x0, y0), self.H, self.W)

    # 可视化地图初始化
    def map_init(self, csv_path="", update_=False, random_flip=0):
        canvas = np.ones(self.canvas_shape, dtype=np.uint8) * 255
        # draw row line
        for row in range(1, self.grid_map.shape[0] + 1):
            cv2.line(canvas, (0, int(self.canvas_shape[0] * row / self.grid_map.shape[0])),
                     (self.canvas_shape[1], int(self.canvas_shape[0] * row / self.grid_map.shape[0])), (255, 0, 0), 1)

        for col in range(1, self.grid_map.shape[1] + 1):
            cv2.line(canvas, (int(self.canvas_shape[1] * col / self.grid_map.shape[1]), 0),
                     (int(self.canvas_shape[1] * col / self.grid_map.shape[1]), self.canvas_shape[0]), (0, 0, 255), 1)

        # 更新气体质量分数可视化地图
        if os.path.exists(csv_path) and update_:
            # 载入真实值的气体分数地图
            self.mass_fraction_viewer.merge(csv_path)
            if random_flip == 0:
                self.mass_fraction_map = self.mass_fraction_viewer.init_map
            elif random_flip == 1:
                self.mass_fraction_map = np.flip(self.mass_fraction_viewer.init_map)
            elif random_flip == 2:
                self.mass_fraction_map = np.flipud(self.mass_fraction_viewer.init_map)
            elif random_flip == 3:
                self.mass_fraction_map = np.fliplr(self.mass_fraction_viewer.init_map)
            # 载入模拟真实情况的高斯噪声地图
            self.Gaussian_mass_map = self.mass_fraction_viewer.noisy_map
            # Blend canvas and mass_fraction_map
            alpha = 0.3  # Weight for canvas
            beta = 0.7  # Weight for mass_fraction
            gamma = 0.5  # Brightness adjustment
            self.obstacle_map = self.obstacle_additive()
            # 融合气体质量分数地图与栅格地图
            obstacle_map = cv2.resize(self.obstacle_map, self.canvas_shape[1::-1])
            obstacle_map_normalized = cv2.normalize(obstacle_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            kernel = np.ones((5, 5), np.uint8)
            self.obstacle_map_normalized = cv2.erode(obstacle_map_normalized, kernel, iterations=5)
            # cv2.imshow('1', obstacle_map_normalized)
            # cv2.waitKey(0)
            mass_fraction_map = cv2.resize(self.mass_fraction_map, self.canvas_shape[1::-1])
            mass_fraction = cv2.normalize(mass_fraction_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
            mass_fraction = cv2.applyColorMap(mass_fraction, cv2.COLORMAP_JET)
            canvas = cv2.addWeighted(canvas, alpha, mass_fraction, beta,
                                     gamma)
            canvas = cv2.bitwise_not(canvas, canvas, mask=self.obstacle_map_normalized)

        self.canvas = canvas

    # 添加障碍层
    def obstacle_additive(self):
        output_map = np.zeros((self.H, self.W))
        if self.obstacle_list is not None:
            for x1, y1, x2, y2 in self.obstacle_list:
                output_map = self.add_Obstacle((x1, y1), (x2, y2), output_map)
        # 构建障碍物掩膜
        return output_map

    # 多线程展现画布
    def show(self):
        while True:
            cv2.imshow(self.title, self.canvas)
            cv2.waitKey(1)


class Point(Colors):
    def __init__(self, index, name, x_max, x_min, y_max, y_min, total_stride, icon_shape=(50, 50)):
        super(Point, self).__init__()
        self.x = 0
        self.y = 0
        self.x_min = x_min
        self.x_max = x_max
        self.y_min = y_min
        self.y_max = y_max
        self.index = index
        self.name = name
        # 记录开始位置
        self.start_point = None
        self.current_grid_pos = None  # x,y
        self.current_pixel_pos = None  # x,y
        self.icon = np.full((*icon_shape, 3), self.rgb_call(index), dtype=np.uint8)
        self.pixel_list = []  # 记录机器走过的路径，像素坐标(x,y)
        self.grid_pos_list = []  # 记录机器走过的路径，栅格坐标(x,y)
        # 初始总步长
        self.total_stride = total_stride
        self.ep_return = 0  # 回归经验值
        self.stride_left = self.total_stride  # 训练步长

    def set_position(self, x, y, start_point):
        self.x = self.clamp(x, self.x_min, self.x_max - self.icon.shape[1])
        self.y = self.clamp(y, self.y_min, self.y_max - self.icon.shape[0])
        # 初始化，添加当前位置
        self.current_pixel_pos = (self.x, self.y)
        self.current_grid_pos = start_point
        # 初始化，添加当前位置
        self.pixel_list.append((self.x, self.y))
        self.grid_pos_list.append(self.current_grid_pos)

    def move(self, del_x, del_y):
        next_x = del_x + self.x
        next_y = del_y + self.y
        # 如果移动之后的距离超出了，返回False表明这次的移动是失败的
        if self.x_min < next_x < self.x_max - self.icon.shape[1] and self.y_min < next_y < self.y_max - self.icon.shape[
            0]:
            self.x += del_x
            self.y += del_y
            # 更新当前坐标
            self.current_pixel_pos = (self.x, self.y)
            return True
        else:
            return False

    def move_pos(self, now_x, now_y):
        next_x = now_x
        next_y = now_y
        # 如果移动之后的距离超出了，返回False表明这次的移动是失败的
        if self.x_min < next_x < self.x_max - self.icon.shape[1] and self.y_min < next_y < self.y_max - self.icon.shape[
            0]:
            self.x = now_x
            self.y = now_y
            # 更新当前坐标
            self.current_pixel_pos = (self.x, self.y)
            return True
        else:
            return False

    def clamp(self, n, minn, maxn):
        return max(min(maxn, n), minn)


class RRT_planner:
    def __init__(self, obstacle_map, step_size=2):
        self.obstacle_map = obstacle_map
        self.step_size = step_size

    @staticmethod
    def bresenham_line(x0, y0, x1, y1):
        """Generate grid cells between two points using Bresenham's line algorithm."""
        points = []
        dx = abs(x1 - x0)
        dy = abs(y1 - y0)
        sx = 1 if x0 < x1 else -1
        sy = 1 if y0 < y1 else -1
        err = dx - dy

        while True:
            points.append((x0, y0))
            if x0 == x1 and y0 == y1:
                break
            e2 = 2 * err
            if e2 > -dy:
                err -= dy
                x0 += sx
            if e2 < dx:
                err += dx
                y0 += sy
        return points

    def is_path_free(self, start, end):
        """Check if the path between start and end is obstacle-free."""
        points = self.bresenham_line(start[0], start[1], end[0], end[1])
        for p in points:
            if self.obstacle_map[p[1], p[0]] == 1:  # Obstacle detected
                return False
        return True

    def nearest_free_grid(self, point):
        """Map a continuous point to the nearest free grid cell using BFS."""
        y, x = int(point[1]), int(point[0])
        if 0 <= y < self.obstacle_map.shape[0] and 0 <= x < self.obstacle_map.shape[1] and self.obstacle_map[y, x] == 0:
            return (x, y)
        queue = deque([(x, y)])
        visited = set([(x, y)])
        directions = [(0, 1), (0, -1), (1, 0), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1)]
        while queue:
            cx, cy = queue.popleft()
            for dx, dy in directions:
                nx, ny = cx + dx, cy + dy
                if 0 <= nx < self.obstacle_map.shape[1] and 0 <= ny < self.obstacle_map.shape[0]:
                    if self.obstacle_map[ny, nx] == 0:
                        return nx, ny
                    if (nx, ny) not in visited:
                        visited.add((nx, ny))
                        queue.append((nx, ny))
        return None

    def rrt_planner(self, start_point, end_point, max_iterations=10000):
        """RRT planner to find a path from start_point to end_point."""
        tree = {start_point: None}  # Tree stores node: parent pairs
        for _ in range(max_iterations):
            # Sample a point (10% chance to sample the endpoint)
            if np.random.random() < 0.1:
                sample = end_point
            else:
                sample = (np.random.randint(0, self.obstacle_map.shape[1] - 1),
                          np.random.randint(0, self.obstacle_map.shape[0] - 1))

            # Find nearest node in the tree
            min_dist = float('inf')
            nearest = None
            for node in tree:
                dist = np.hypot(node[0] - sample[0], node[1] - sample[1])
                if dist < min_dist:
                    min_dist = dist
                    nearest = node

            # Extend towards the sampled point
            direction = np.array(sample) - np.array(nearest)
            length = np.linalg.norm(direction)
            if length > self.step_size:
                direction = direction / length * self.step_size
            new_point = np.array(nearest) + direction
            new_point = (int(new_point[0]), int(new_point[1]))

            # Map to the nearest free grid，检测这个新的点是否会发生碰撞
            new_point = self.nearest_free_grid(new_point)
            if new_point is None:
                continue

            # Check if the path is free and add to tree
            if self.is_path_free(nearest, new_point):
                tree[new_point] = nearest
                # Check if close to the endpoint,计算欧式距离查看是否离的够近
                if np.hypot(new_point[0] - end_point[0], new_point[1] - end_point[1]) <= self.step_size:
                    path = []
                    current = new_point
                    while current is not None:
                        path.append(current)
                        if current == tree[current]:
                            path = []
                            break
                        current = tree[current]  # value -> key
                    path.reverse()
                    return path
        return []  # No path found


# 终端用于控制每个机器人，便于生成多个机器人
class Terminal(Map, RRT_planner):
    def __init__(self, canvas_size=(800, 800), split_x_nums=20, split_y_nums=20, obstacle_list=None):
        Map.__init__(self, canvas_size=canvas_size, split_x_nums=split_x_nums, split_y_nums=split_y_nums,
                     obstacle_list=obstacle_list)
        RRT_planner.__init__(self, self.obstacle_map)
        # 初始化地图
        self.map_init()
        # 通过一个字典管理所有机器人
        self.element_collection = None
        # 离散化移动的方式以及步骤
        self.box_min_step = np.array([int(self.canvas_shape[1] / self.grid_map.shape[1]),
                                      int(self.canvas_shape[0] / self.grid_map.shape[0])])
        self.action_dict = {"upper": self.box_min_step * np.array([0, -1]),
                            "lower": self.box_min_step * np.array([0, 1]),
                            "left": self.box_min_step * np.array([-1, 0]),
                            "right": self.box_min_step * np.array([1, 0]),
                            "up_right": self.box_min_step * np.array([1, -1]),
                            "up_left": self.box_min_step * np.array([-1, -1]),
                            "low_right": self.box_min_step * np.array([1, 1]),
                            "low_left": self.box_min_step * np.array([-1, 1])}
        self.path_dict = {}

    def RRT_local_planner(self, start_point, stride=5):
        for key in self.action_dict:
            next_pos = start_point
            for _ in range(stride):
                next_step = next_pos + np.array(self.action_dict[key] / self.box_min_step, dtype=np.int32)
                if 0 <= next_step[0] < self.grid_map.shape[1] and 0 <= next_step[1] < self.grid_map.shape[0]:
                    next_pos = next_step
            self.path_dict[key] = self.rrt_planner(start_point, next_pos)

    # 将所有元素呈现在画布上
    def draw_elements_on_canvas(self, csv_path, random_flip=0):
        # Init the canvas
        self.map_init(csv_path=csv_path, update_=True, random_flip=random_flip)
        # 将每个机器人画出来
        for index, Bot_name in enumerate(self.element_collection.keys()):
            text = f"Bot-{Bot_name} | Stride Left: {self.element_collection[Bot_name].stride_left} | Rewards: {self.element_collection[Bot_name].ep_return}"
            # Put the info on canvas
            cv2.putText(self.canvas, text, (10, 20 * (index + 1)), font,
                        1, (0, 0, 0), 1, cv2.LINE_AA)
            # Draw the heliopter on canvas
            x, y = self.element_collection[Bot_name].x, self.element_collection[Bot_name].y
            elem_shape = self.element_collection[Bot_name].icon.shape
            self.canvas[y: y + elem_shape[0], x:x + elem_shape[1]] = self.element_collection[Bot_name].icon
            # for index, point in enumerate(self.element_collection[Bot_name].pixel_list):
            #     if index < len(self.element_collection[Bot_name].pixel_list) - 1:
            #         cv2.line(self.canvas, self.element_collection[Bot_name].pixel_list[index], self.element_collection[Bot_name].pixel_list[index + 1], (0, 255, 0), 5)
            #     cv2.circle(self.canvas, point, 5, (0, 255, 0), -1)

    def create_bot(self, bot_elements: dict):
        element_collection = {}
        icon_shape = (
            int((self.canvas_shape[1] / self.grid_map.shape[1]) * 0.3),
            int((self.canvas_shape[0] / self.grid_map.shape[0]) * 0.3))

        origin_step = (int(self.canvas_shape[1] / (2 * self.grid_map.shape[1]) - icon_shape[1] / 2),
                       int(self.canvas_shape[0] / (2 * self.grid_map.shape[0]) - icon_shape[0] / 2))

        # 统一处理字典或列表
        elements = bot_elements.values() if isinstance(bot_elements, dict) else bot_elements
        for index, element in enumerate(elements):
            # 实例化机器人本体
            element_collection[element.name] = Point(index, name=element.name, x_max=self.x_max, x_min=self.x_min,
                                                     y_max=self.y_max, y_min=self.y_min,
                                                     total_stride=element.total_stride,
                                                     icon_shape=icon_shape)
            # 设定初始位置
            start_point = element.start_point
            if start_point[0] < self.grid_map.shape[1] and start_point[1] < self.grid_map.shape[0]:
                x_pos = int(start_point[0] * self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0])
                y_pos = int(start_point[1] * self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1])
                self.grid_map[start_point[1], start_point[0]] = 1  # 更新栅格地图
            else:
                x_pos = int(self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0])
                y_pos = int(self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1])

            # 初始化，设置所有点坐标
            element_collection[element.name].set_position(x_pos, y_pos, start_point)

        return element_collection

    # 更新机器人在栅格地图中位置的函数
    def upgrade_grid_position(self, bot_name):
        current_state = self.element_collection[bot_name].current_pixel_pos
        current_state = int(current_state[0] * self.grid_map.shape[1] / self.canvas_shape[1]), int(
            current_state[1] * self.grid_map.shape[0] / self.canvas_shape[0])
        # 记录已经走过的路径
        if self.obstacle_map is not None and not self.obstacle_map[current_state[1], current_state[0]]:
            self.grid_map[current_state[1], current_state[0]] = 1  # 在栅格地图上赋予机器人位置信息
            self.element_collection[bot_name].current_grid_pos = current_state  # 为机器人附上栅格地图的位置
            return True
        else:
            return False

    def take_actions_pos(self, pos_step, bot_name):
        assert self.element_collection[bot_name] is not None, "Please add element to execute following task"

        # 将栅格点转换成像素点，进行移动
        def grid_2_pixel():
            icon_shape = (
                int((self.canvas_shape[1] / self.grid_map.shape[1]) * 0.3),
                int((self.canvas_shape[0] / self.grid_map.shape[0]) * 0.3))
            origin_step = (int(self.canvas_shape[1] / (2 * self.grid_map.shape[1]) - icon_shape[1] / 2),
                           int(self.canvas_shape[0] / (2 * self.grid_map.shape[0]) - icon_shape[0] / 2))
            if pos_step[0] < self.grid_map.shape[1] and pos_step[1] < self.grid_map.shape[0]:
                x_pos = int(pos_step[0] * self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0])
                y_pos = int(pos_step[1] * self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1])
            else:
                x_pos = int(self.canvas_shape[1] / self.grid_map.shape[1] + origin_step[0])
                y_pos = int(self.canvas_shape[0] / self.grid_map.shape[0] + origin_step[1])
            return x_pos, y_pos

        move_flag = self.element_collection[bot_name].move_pos(*grid_2_pixel())
        upgrade_flag = self.upgrade_grid_position(bot_name)
        # 移动机器人，减少可移动步数，如果是假设移动，则不减少步数
        self.element_collection[bot_name].stride_left -= 1
        if move_flag and upgrade_flag:
            # 将机器人走的路径添加到坐标里
            self.element_collection[bot_name].pixel_list.append(self.element_collection[bot_name].current_pixel_pos)
            self.element_collection[bot_name].grid_pos_list.append(
                self.element_collection[bot_name].current_grid_pos)

        # 返回布尔值，确定这一步是否是正常的
        return True if upgrade_flag and move_flag else False

    def take_actions(self, action, bot_name):  # 给予一个sim_仿真符号，为了测试运动是否可行的同时，不更改机器人动作
        if not isinstance(action, str):  # 如果是str命令
            action = list(self.action_dict.keys())[action]
        assert action in self.action_dict.keys(), "Invalid Action"
        move_flag = self.element_collection[bot_name].move(self.action_dict[action][0],
                                                           self.action_dict[action][1])

        upgrade_flag = self.upgrade_grid_position(bot_name)
        # 移动机器人，减少可移动步数，如果是假设移动，则不减少步数
        self.element_collection[bot_name].stride_left -= 1
        if move_flag and upgrade_flag:
            # 将机器人走的路径添加到坐标里
            self.element_collection[bot_name].pixel_list.append(self.element_collection[bot_name].current_pixel_pos)
            self.element_collection[bot_name].grid_pos_list.append(
                self.element_collection[bot_name].current_grid_pos)

        # 返回布尔值，确定这一步是否是正常的
        return True if upgrade_flag and move_flag else False

    def _get_state(self, bot_name, action):
        # 当前位置
        next_pos = self.element_collection[bot_name].current_grid_pos
        # 1.确定当前位置周围3*3的气体质量分数
        max_mass_thresh = np.max(self.mass_fraction_map)
        fraction_grid = []
        for key in self.action_dict:
            next_step = next_pos + np.array(self.action_dict[key] / self.box_min_step, dtype=np.int32)
            # 在地图范围内赋予气体值，否则赋予-1
            if 0 <= next_step[0] < self.grid_map.shape[1] and 0 <= next_step[1] < self.grid_map.shape[0] and \
                    not self.obstacle_map[next_step[1], next_step[0]]:
                fraction_grid.append(self.mass_fraction_map[tuple(next_step[::-1])] / max_mass_thresh)
            else:
                fraction_grid.append(-1)

        # 2.确定当前位置八个方向向前最终到达的障碍物距离
        obstacle_grid = []
        for key in self.action_dict:
            next_step = next_pos + np.array(self.action_dict[key] / self.box_min_step, dtype=np.int32)
            distance_2_obstacle = 0
            while True:
                if 0 <= next_step[0] < self.grid_map.shape[1] and 0 <= next_step[1] < self.grid_map.shape[0] and \
                        not self.obstacle_map[next_step[1], next_step[0]]:
                    next_step = next_step + np.array(self.action_dict[key] / self.box_min_step, dtype=np.int32)
                    distance_2_obstacle += 1
                else:
                    break
            obstacle_grid.append(distance_2_obstacle)
        obstacle_grid = np.array(obstacle_grid) / np.sqrt(self.W ** 2 + self.H ** 2)
        obstacle_grid = np.where(obstacle_grid == 0, -1, obstacle_grid)

        # 3.确定动作方向维度的状态
        action_dim = np.zeros(len(self.action_dict))
        action_dim[list(self.action_dict.keys()).index(action)] = 1

        # 4.机器人自身位置，处于全局的位置
        next_pos = np.array(next_pos) / np.array([self.W, self.H])
        return np.array(
            fraction_grid + action_dim.tolist() + obstacle_grid.tolist() + next_pos.tolist(),
            dtype=np.float64)

    def step(self, action, bot_name,pos_step=None, view=True):
        if not isinstance(action, str):
            action = list(self.action_dict.keys())[action]
        assert self.element_collection[bot_name] is not None, "Please add element to execute following task"
        assert action in self.action_dict.keys(), "Invalid Action"
        epoch_reward, done_s = 0, False
        # 移动前的位置和统计
        before_point_sum = self.grid_map.sum()
        # 移动前的位置
        current_pos = self.element_collection[bot_name].current_grid_pos
        # 1.先进行移动take a action
        if not self.take_actions(action, bot_name=bot_name):
            print("knocked wall!!!")
            # epoch_reward = int(-self.element_collection[bot_name].ep_return - 10)
            epoch_reward = -5
            # done_s = True
        # 移动后的位置
        next_pos = self.element_collection[bot_name].current_grid_pos
        # 移动后的位置和统计
        after_point_sum = self.grid_map.sum()
        # 移动后的位置的气体质量分数
        next_this_point_mass = self.mass_fraction_map[next_pos[1], next_pos[0]]

        # 2.转换为ndarray数组输出的地图，表示质量分数地图
        # # 执行膨胀操作
        dilated_path = cv2.dilate(self.grid_map, self.kernel, iterations=1)
        quality_scores_map = np.array(dilated_path * self.mass_fraction_map)
        # # 将quality_scores_map放缩到[0,1]

        k1, k2, k3, k4 = 0.5, 0.2, 2.5, 0.7  # 设置权重影响分数·

        # 2.获取奖励
        def _policy_():
            episode_reward = 0  # 初始化每轮的奖励
            # 1.Dynamic Exploration Reward
            stride_ratio = self.element_collection[bot_name].stride_left / self.element_collection[
                bot_name].total_stride
            visited_ratio = self.grid_map.sum() / (self.grid_map.shape[0] * self.grid_map.shape[1])  # 时间因素
            if after_point_sum != before_point_sum:  # 鼓励探索，但是惩罚来回震荡
                visited_reward = k1 * (1 - visited_ratio)
                exploration_bonus = stride_ratio * visited_reward
            else:
                visited_reward = k2 * (1 - visited_ratio)
                exploration_bonus = - visited_reward
            # print(exploration_bonus, "exploration_bonus")
            # 2.Max closer reward bonus,选择最大的奖励，加入自身位置，是防止来回刷分
            Max_thresh_local = 0
            Max_pos_choice = current_pos

            for key in self.action_dict:
                next_step = current_pos + np.array(self.action_dict[key] / self.box_min_step, dtype=np.int32)
                # 在地图范围内赋予气体值，否则赋予-1
                if 0 <= next_step[0] < self.grid_map.shape[1] and 0 <= next_step[1] < self.grid_map.shape[0] and \
                        not self.obstacle_map[tuple(next_step[::-1])]:

                    # print(next_step, self.mass_fraction_map[tuple(next_step[::-1])])
                    if Max_thresh_local < self.mass_fraction_map[tuple(next_step[::-1])]:
                        Max_thresh_local = self.mass_fraction_map[tuple(next_step[::-1])]
                        Max_pos_choice = next_step

            Max_choice_bonus = k3 * (
                    self.mass_fraction_map[tuple(Max_pos_choice[::-1])] / np.max(self.mass_fraction_map))
            Max_choice_bonus = Max_choice_bonus if np.linalg.norm(
                np.array(next_pos) - np.array(Max_pos_choice)) <= 1 else -0.2

            # print(next_pos, 'next_pos')
            # print(Max_pos_choice, 'Max_pos_choice')
            # print(Max_choice_bonus, "Max_choice_reward")
            episode_reward += exploration_bonus + \
                              Max_choice_bonus + \
                              - 0.1

            # print(episode_reward, 'episode_reward')
            return episode_reward

        epoch_reward += _policy_() if not done_s else 0

        if next_this_point_mass > k4 * np.max(self.mass_fraction_map):
            epoch_reward = 50
            done_s = True
        self.element_collection[bot_name].ep_return += epoch_reward

        if self.element_collection[bot_name].stride_left < 1 or done_s:
            print('\n', "本轮得分:", '\n', self.element_collection[bot_name].ep_return)
            if view:
                self.fit_map()
                plt.clf()
                plt.imshow(quality_scores_map, cmap='jet', vmin=0,
                           vmax=np.max(quality_scores_map))  # cmap 指定颜色映射
                for direction, path in self.path_dict.items():
                    if path:  # Check if the path is not empty
                        x, y = zip(*path)  # Separate x and y coordinates
                        plt.plot(x, y, marker='o', label=direction)  # Plot path with points and lines

                plt.colorbar(label='Value')  # 添加颜色条
                plt.title('quality_scores_map_ Example')  # 设置标题
                plt.xlabel('X-axis')  # X 轴标签
                plt.ylabel('Y-axis')  # Y 轴标签
                plt.show()
            done_s = True  # 表明状态已经完成

        # 3.获取状态

        state = self._get_state(bot_name, action)
        # print("单步输入状态", state)
        return state, epoch_reward, done_s

    # 设置控制模式

    def render(self, bot_elements, mode="parallel", ):
        assert mode in ["parallel", "step"], "Invalid mode, must be either \"parallel\" or \"step\""
        # assert self.mass_fraction_map is not None
        self.element_collection = self.create_bot(bot_elements=bot_elements)
        if mode == "parallel":
            # 渲染的时候，直接开个新的线程进行并行渲染。
            t1 = threading.Thread(target=self.show, args=(), daemon=True)
            # 设置守护线程,即主线程启动才会启动子线程
            t1.start()
        elif mode == "step":
            pass

    def reset(self, bot_elements):
        # 重置每一个机器人的得分
        state_list = []
        self.element_collection = self.create_bot(bot_elements)
        self.grid_map = np.zeros(self.grid_map.shape)  # 栅格概率地图重置
        for key in bot_elements:
            state_list.append(self._get_state(bot_elements[key].name, "lower"))
        return state_list


if __name__ == "__main__":
    # 创造一些机器人对象，便于获取
    Bot_elements = {"Bot1": BOT("Bot1", (15, 15), 1)}
    terminal = Terminal(canvas_size=(900, 600), split_x_nums=30, split_y_nums=45,
                        obstacle_list=[(0, 26, 7, 24), (20, 17, 26, 15), (14, 7, 15, 3)])
    terminal.render(Bot_elements)
    terminal.map_init(csv_path=r'D:\Gas_detector\data\gas_mass_data\export15.csv', update_=True)
    while True:
        terminal.draw_elements_on_canvas(csv_path=r'D:\Gas_detector\data\gas_mass_data\export15.csv')
        cv2.imshow("mass", cv2.resize(terminal.mass_fraction_map, (600, 900)))
        key = cv2.waitKey(60)
        for bot in terminal.element_collection.values():
            if key == 119:
                terminal.step('upper', f"{bot.name}")
            elif key == 115:
                terminal.step('lower', f"{bot.name}")
            elif key == 100:
                terminal.step('right', f"{bot.name}")
            elif key == 97:
                terminal.step('left', f"{bot.name}")
            elif key == 113:
                print(terminal.reset(Bot_elements))
                continue
